#ifndef __MULTIGRID_H__
#define __MULTIGRID_H__

#include <vector>
#include <lapacke.h>
#include "MultiArrays.h"
#include "BoundaryConditions.h"
#include "InterGridOperator.h"
#include "norm.h"

using RealIter = std::vector<Real>::iterator;

inline void Weighted_Jacobi_Relax(Real Mat[], int size_N, RealIter x_begin, RealIter b_begin, double w, int v);

inline void Compute_Residue(Real Mat[], int size_N, RealIter x_begin, RealIter b_begin);

template <int N>
class MultiGridSolver
{
public:
    // if direct solve the bottom equations
    bool isDirectBottomSolver = false;

    const GridOp::Interpolation<Real, N> *_interpolate;
    const GridOp::Restriction<Real, N> *_restrict;

    MultiGridSolver(const GridOp::Interpolation<Real, N> *interpolate, const GridOp::Restriction<Real, N> *restrict) : _interpolate(interpolate), _restrict(restrict) {}
    MultiGridSolver(const GridOp::Interpolation<Real, N> &interpolate, const GridOp::Restriction<Real, N> &restrict) : _interpolate(&interpolate), _restrict(&restrict) {}
    ~MultiGridSolver();

    void assignF(std::function<Real(Point<N>)> &&F);
    void assignF(const std::function<Real(Point<N>)> &F);

    void define_boundary(BoundaryConditions<N> &&BC) { _BC = BC; }

    /**
     * @brief discretize the domain into multigrids
     * and in this subroutine the matrices and the
     * right hand side will be filled.
     * Make sure this function is called after
     * 'assignF' and 'define_boundary' and befor 'solve'.
     *
     * @param n
     */
    void discretization(int n);

    /**
     * @brief
     *
     * @param solution : the solution vector
     * @param method : "FMG" , "VCycle" , "LU"
     * @param v1 : presmoothing
     * @param v2 : postsmoothing
     * @param epsil : the cretiria ending iteration
     * @param M : the Maximum iterations
     */
    void solve(std::vector<Real> &solution, const std::string &method, int v1 = 2, int v2 = 2, Real epsil = 1e-8, int M = 20);

    void PrintAdj(std::ostream &out);

    void PrintMatrix(std::ostream &out);

    const std::vector<Point<N>> &getGrid() const { return grid_coords; }

    /**
     * @brief Get the 2-norm of residuals
     *
     * @return Real
     */
    Real getResidNorm();

private:
    // -\nabla u = f
    std::function<Real(Point<N>)> _F;

    BoundaryConditions<N> _BC;

    //_v : the approximation of solutions
    //_f : the right hand side of equations
    MultiArrays<Real, N> _v, _f;
    //_A : the coefficient matrices, each pointer pointing to one single row in the matrix
    MultiArrays<Real *, N> _A;
    MultiArrays<int *, N> adjacency;

    std::vector<Point<N>> grid_coords;

    void VCycle(std::vector<Real> &solution, int v1, int v2);
    void VCycle_helper(int v1, int v2, int l);

    void FMG(std::vector<Real> &solution, int v1, int v2);
    void FMG_helper(int v1, int v2, int l);

    // internal iterative helper function to add point to the grids
    void add_point(int dim, Point<N> &p, Real value, int n);

    // fill the matrix
    void construct_mat(int n);

    /**
     * @brief fill the row coresponding to a boundary point
     *
     * @param row
     * @param bound_idx : the single boundary the point located
     * @param point_idx
     * @param l : grid level
     */
    inline void fill_boundary_row(Real *row, int bound_idx, int point_idx, int l);
};

template <int N>
MultiGridSolver<N>::~MultiGridSolver()
{
    for (int i = 0; i < _A.getLevelNum(); ++i)
    {
        delete[](*_A.Level(i));
        delete[](*adjacency.Level(i));
    }
}

template <int N>
void MultiGridSolver<N>::assignF(std::function<Real(Point<N>)> &&F)
{
    _F = F;
    for (int i = 0; i < grid_coords.size(); ++i)
    {
        if (_BC.Location(grid_coords[i]) > 0)
            _f.Level(0)[i] = _F(grid_coords[i]);
    }
}

template <int N>
void MultiGridSolver<N>::assignF(const std::function<Real(Point<N>)> &F)
{
    _F = F;
    for (int i = 0; i < grid_coords.size(); ++i)
    {
        if (_BC.Location(grid_coords[i]) > 0)
            _f.Level(0)[i] = _F(grid_coords[i]);
    }
}
template <int N>
void MultiGridSolver<N>::solve(std::vector<Real> &solution, const std::string &method, int v1, int v2, Real epsil, int M)
{
    if (method == "VCycle")
    {
        double old_norm = getResidNorm(), new_norm = old_norm;
        for (int i = 0; i < M; ++i)
        {
            VCycle(solution, v1, v2);
            old_norm = new_norm;
            new_norm = getResidNorm();
            double factor = new_norm / old_norm;
            std::cout << "VCycle iteration : " << i + 1 << " , redidual norm : " << new_norm << " , reduction factor : " << factor << std::endl;
            if (new_norm < epsil || (factor > 0.95 && new_norm < 1e-6))
                break;
        }
    }
    else if (method == "FMG")
    {
        FMG(solution, v1, v2);
    }
    else if (method == "LU")
    {
        int size_N = _v.getLevelPointNum(0);
        int nrhs = 1, INFO;
        std::vector<int> IPIV(size_N);
        std::vector<double> b(size_N);
        std::vector<Real> A_copy(size_N * size_N);
        for (int i = 0; i < size_N * size_N; ++i)
        {
            A_copy[i] = (*_A.Level(0))[i];
        }

        std::copy(_f.Level(0), _f.Level(1), b.begin());

        dgesv_(&size_N, &nrhs, A_copy.data(), &size_N, IPIV.data(), b.data(), &size_N, &INFO);
        if (INFO == 0)
        {
            std::copy(b.begin(), b.end(), solution.begin());
        }
        else if (INFO > 0)
            std::cout << " lapack_dgesv error: A is singular " << std::endl;
        else
            std::cout << " lapack_dgesv: the " << -INFO << "-th argument has an illegal value" << std::endl;

        return;
    }
    else
    {
        std::cout << "Please input a valid method!" << std::endl;
    }
}

template <int N>
void MultiGridSolver<N>::add_point(int dim, Point<N> &p, Real value, int n)
{
    p[dim] = value;
    if (dim == 0)
        grid_coords.push_back(p);
    else
    {
        Real h = 1.0 / n;
        for (int k = 0; k <= n; ++k)
        {
            add_point(dim - 1, p, k * h, n);
        }
    }
}

template <int N>
void MultiGridSolver<N>::discretization(int n)
{
    if (_A.getLevelNum() > 0)
    {
        for (int i = 0; i < _A.getLevelNum(); ++i)
        {
            delete[](*_A.Level(i));
            delete[](*adjacency.Level(i));
        }
    }
    _v.resize(n);
    _f.resize(n);
    _A.resize(n);
    adjacency.resize(n);
    double h = 1.0 / n;

    grid_coords = std::vector<Point<N>>();
    Point<N> tmp;
    for (int k = 0; k <= n; ++k)
        add_point(N - 1, tmp, k * h, n);

    // initialize all the multiarrays
    for (int i = 0, current_size = n; i < _A.getLevelNum(); ++i, current_size /= 2)
    {
        int size = _A.getLevelPointNum(i);
        auto current_mat = _A.Level(i);
        auto adjacency_lists = adjacency.Level(i);
        auto current_v = _v.Level(i);
        auto current_f = _f.Level(i);

        *adjacency_lists = new int[size * (1 << N)];
        for (int j = 0; j < size * (1 << N); ++j)
        {
            (*adjacency_lists)[j] = -1;
        }
        *current_mat = new Real[size * size]{0};

        Real *current_row = *current_mat;
        int *current_adj_list = *adjacency_lists;
        for (int j = 1; j < size; ++j)
        {
            current_row += size;
            current_mat[j] = current_row;
            current_adj_list += (1 << N);
            adjacency_lists[j] = current_adj_list;
        }

        for (int j = 0; j < size; ++j)
        {
            current_v[j] = 0;
            current_f[j] = 0;
        }
    }

    // record the adjcency of each point
    // the two points close to p in dimension dim will be placed in index 2*dim and 2*dim + 1
    for (int i = 0, curent_size = n; i < adjacency.getLevelNum(); ++i, curent_size /= 2)
    {
        int size_N = adjacency.getLevelPointNum(i);
        auto current_adj_list = adjacency.Level(i);
        for (int j = 0; j < size_N; ++j)
        {
            for (int dim = 0, coord_unit = 1; dim < N; ++dim, coord_unit *= (curent_size + 1))
            {
                if (0 <= j - coord_unit && j - coord_unit < size_N)
                {
                    current_adj_list[j][2 * dim] = (j - coord_unit);
                }
                if (0 <= j + coord_unit && j + coord_unit < size_N)
                {
                    current_adj_list[j][2 * dim + 1] = (j + coord_unit);
                }
            }
        }
    }

    construct_mat(n);
}

template <int N>
void MultiGridSolver<N>::VCycle_helper(int v1, int v2, int l)
{
    // solve the bottom equations directly
    if (l == _v.getLevelNum() - 1 && isDirectBottomSolver)
    {
        int size_N = _v.getLevelPointNum(l);
        int nrhs = 1, INFO;
        std::vector<int> IPIV(size_N);
        std::vector<double> b(size_N);
        std::vector<Real> A_copy(size_N * size_N);
        for (int i = 0; i < size_N * size_N; ++i)
        {
            A_copy[i] = (*_A.Level(l))[i];
        }

        std::copy(_f.Level(l), _f.Level(l + 1), b.begin());

        dgesv_(&size_N, &nrhs, A_copy.data(), &size_N, IPIV.data(), b.data(), &size_N, &INFO);
        if (INFO == 0)
        {
            std::copy(b.begin(), b.end(), _v.Level(l));
        }
        else if (INFO > 0)
            std::cout << " lapack_dgesv error: A is singular " << std::endl;
        else
            std::cout << " lapack_dgesv: the " << -INFO << "-th argument has an illegal value" << std::endl;

        return;
    }

    // 1.relaxation
    Weighted_Jacobi_Relax(*_A.Level(l), _v.getLevelPointNum(l), _v.Level(l), _f.Level(l), 2.0 / 3, v1);

    if (l < _v.getLevelNum() - 1)
    {
        std::vector<Real> f_copy(_f.getLevelPointNum(l));
        std::copy(_f.Level(l), _f.Level(l + 1), f_copy.begin());
        Compute_Residue(*_A.Level(l), _v.getLevelPointNum(l), _v.Level(l), _f.Level(l));
        (*this->_restrict)(_f, l);
        std::copy(f_copy.begin(), f_copy.end(), _f.Level(l));

        for (auto it = _v.Level(l + 1); it != _v.Level(l + 2); ++it)
        {
            *it = 0;
        }

        VCycle_helper(v1, v2, l + 1);
        (*this->_interpolate)(_v, l + 1, false);
    }

    // 2.another relaxation
    Weighted_Jacobi_Relax(*_A.Level(l), _v.getLevelPointNum(l), _v.Level(l), _f.Level(l), 2.0 / 3, v2);
}

template <int N>
void MultiGridSolver<N>::VCycle(std::vector<Real> &solution, int v1, int v2)
{
    std::copy(solution.begin(), solution.end(), _v.Level(0));
    VCycle_helper(v1, v2, 0);
    std::copy(_v.Level(0), _v.Level(1), solution.begin());
}

template <int N>
void MultiGridSolver<N>::FMG(std::vector<Real> &solution, int v1, int v2)
{
    FMG_helper(v1, v2, 0);
    std::copy(_v.Level(0), _v.Level(1), solution.begin());
}

template <int N>
void MultiGridSolver<N>::FMG_helper(int v1, int v2, int l)
{
    if (l < _v.getLevelNum() - 1)
    {
        (*this->_restrict)(_f, l);
        (*this->_restrict)(_v, l);
        FMG_helper(v1, v2, l + 1);
        (*this->_interpolate)(_v, l + 1, true);
    }
    VCycle_helper(v1, v2, l);
}

template <int N>
void MultiGridSolver<N>::construct_mat(int n)
{
    for (int i = 0, current_n = _A.getFinestSize(); i < _A.getLevelNum(); ++i, current_n /= 2)
    {
        auto mat = _A.Level(i);
        int size_N = _A.getLevelPointNum(i);
        for (int j = 0; j < size_N; ++j)
        {
            Real *row = mat[j];
            int t_idx = _v.getTrueIdx(j, i);
            const Point<N> &p = grid_coords[t_idx];

            if (_BC.Location(p) > 0) // internal points
            {
                _f.Level(i)[j] = _F(p);
                row[j] = current_n * current_n * (1 << N);
                for (int dim = 0; dim < N; ++dim)
                {
                    row[adjacency.Level(i)[j][2 * dim]] = -current_n * current_n;
                    row[adjacency.Level(i)[j][2 * dim + 1]] = -current_n * current_n;
                }
            }
            else // points on the boundary
            {
                int bound_idx = 0;
                for (int i = 0; i < _BC.bounds.size(); ++i)
                {
                    if (_BC.bounds[i].Location(p) == 0)
                    {
                        bound_idx = i;
                        break;
                    }
                }
                fill_boundary_row(row, bound_idx, j, i);
            }
        }
    }
}

template <int N>
inline void MultiGridSolver<N>::fill_boundary_row(Real *row, int bound_idx, int point_idx, int l)
{
    int current_n = _A.getFinestSize() >> l;
    Point<N> p = grid_coords[point_idx];
    SingleBoundaryCondition<N> &bound = _BC.bounds[bound_idx];

    //\alpha * u
    row[point_idx] += bound.a;
    _f.Level(l)[point_idx] = bound.g(p);

    if (bound.b != 0)
    {
        Vec<Real, N> nor_vec = bound.Normal_Vec(p);
        auto adj_list = adjacency.Level(l);

        for (int dim = 0; dim < N; ++dim)
        {
            int prev_idx = adj_list[point_idx][2 * dim],
                next_idx = adj_list[point_idx][2 * dim + 1];
            if ((prev_idx != -1) && (next_idx != -1))
            {
                row[prev_idx] += -bound.b * nor_vec[dim] * current_n / 2;
                row[next_idx] += bound.b * nor_vec[dim] * current_n / 2;
            }
            else if ((prev_idx != -1) && (next_idx == -1))
            {
                int prev_prev_idx = adj_list[prev_idx][2 * dim];
                row[point_idx] += bound.b * nor_vec[dim] * 3 * current_n / 2;
                row[prev_idx] += -bound.b * nor_vec[dim] * 2 * current_n;
                row[prev_prev_idx] += bound.b * nor_vec[dim] * current_n / 2;
            }
            else if ((prev_idx == -1) && (next_idx != -1))
            {
                int next_next_idx = adj_list[next_idx][2 * dim + 1];
                row[point_idx] += -bound.b * nor_vec[dim] * 3 * current_n / 2;
                row[next_idx] += bound.b * nor_vec[dim] * 2 * current_n;
                row[next_next_idx] += -bound.b * nor_vec[dim] * current_n / 2;
            }
        }
    }
}

inline void Weighted_Jacobi_Relax(Real Mat[], int size_N, RealIter x, RealIter b, double w, int v)
{
    for (int k = 0; k < v; ++k)
    {
        for (int i = 0; i < size_N; ++i)
        {
            Real d = Mat[i + size_N * i];
            Real x_WJ = 0;
            for (int j = 0; j < size_N; ++j)
            {
                x_WJ += Mat[i * size_N + j] * x[j];
            }
            x[i] = x[i] - w * x_WJ / d + w * b[i] / d;
        }
    }
}

inline void Compute_Residue(Real Mat[], int size_N, RealIter x_begin, RealIter b_begin)
{
    for (int i = 0; i < size_N; ++i)
    {
        for (int j = 0; j < size_N; ++j)
        {
            b_begin[i] -= Mat[i * size_N + j] * x_begin[j];
        }
    }
}

template <int N>
void MultiGridSolver<N>::PrintAdj(std::ostream &out)
{
    for (int l = 0; l < adjacency.getLevelNum(); ++l)
    {
        out << "Level [" << l << "] : " << std::endl;
        auto adj_lists = adjacency.Level(l);
        for (int i = 0; i < adjacency.getLevelPointNum(l); ++i)
        {
            out << grid_coords[adjacency.getTrueIdx(i, l)] << " : ";
            for (int j = 0; j < (1 << N); ++j)
            {
                if (adj_lists[i][j] != -1)
                    out << grid_coords[adjacency.getTrueIdx(adj_lists[i][j], l)] << ",";
            }
            out << std::endl;
        }
    }
}

template <int N>
void MultiGridSolver<N>::PrintMatrix(std::ostream &out)
{
    for (int l = 0; l < _A.getLevelNum(); ++l)
    {
        out << "Level [" << l << "] : " << std::endl;
        auto Mat = _A.Level(l);
        int size_N = _A.getLevelPointNum(l);
        for (int i = 0; i < size_N; ++i)
        {
            for (int j = 0; j < size_N; ++j)
            {
                printf("%.4f,", Mat[i][j]);
            }
            printf("\n");
        }
    }
}

template <int N>
Real MultiGridSolver<N>::getResidNorm()
{
    int size_N = _f.getLevelPointNum(0);
    std::vector<Real> f_copy(size_N);
    std::copy(_f.Level(0), _f.Level(1), f_copy.begin());
    Compute_Residue(*_A.Level(0), size_N, _v.Level(0), _f.Level(0));
    Real norm = my_vec_norm(&(*_f.Level(0)), size_N, 2);
    std::copy(f_copy.begin(), f_copy.end(), _f.Level(0));
    return norm;
}

#endif // __MULTIGRID_H__